#!/usr/bin/env python

# Generic tree abstraction of the results of an LDAP query. The result
# tree is constructed by dividing the search results according to "groups",
# which correspond to organizational units within the LDAP database.
#
# Depends on: python-ldap
#
#
# Author:   Mike Murphy (mamurph@cs.clemson.edu)
# Revision: 14 July 2008
#
#   Copyright 2007-2008 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.

'''Provides a generic tree abstraction of LDAP search results.'''

import ldap


class LDAPRecord(object):
   '''Generic LDAP query result. Contains the DN and verbatim attributes
   resulting from the search. LDAPRecords compare to one another by using
   string comparisons on their respective DN's.'''

   def __init__(self, dn, attributes):
      '''Class constructor: dn (str) is the Distinguished Name of the
      LDAP object. attributes (dict) is the verbatim attributes dictionary
      returned by the LDAP query.'''

      self.dn = dn
      self.attributes = attributes

   def __repr__(self, indent = 0):
      '''String representation of the LDAPRecord, optionally indented by
      a selectable number of spaces.'''

      result = (' ' * indent) + (self.dn.split(',')[0].rstrip(',')) + '\n'
      for attr in self.attributes:
         result += (' ' * indent) + ' * ' + attr + ': ' + \
            str(self.attributes[attr]) + '\n'
      return result

   def __cmp__(self, other):
      '''LDAPRecords compare to each other by comparison of their DN's'''

      # Unfortunately, str does not implement __cmp__, so the following
      # hack is needed
      retcode = 0
      if self.dn < other.dn:
         retcode = -1
      elif self.dn > other.dn:
         retcode = 1
      else:
         retcode = 0
      return retcode


class LDAPTree(object):
   '''LDAPTree objects form the inner (non-leaf) nodes of the results tree.
   Both the global root node, and the root node of each subtree, are of this
   type, and all tree methods can be applied both to the root node and to
   the subtrees. Trees may be changed in place by means of pruning (removing
   a leaf or subtree) or flattening (recursively moving leaves from a child
   subtree up to the parent (sub)tree). The root of the tree stores the base
   DN of the tree, which is the DN used as the search base in the LDAP
   query that produces the tree.'''

   def __init__(self, name, dn):
      '''Constructor (name, dn): name is the "group" name (first segment of the
      DN with ou= removed, or the DN in the case of the root node). dn is the
      Distinguished Name of the organizational unit object corresponding to the
      group name, or the base search DN in the case of the root node.'''

      self.name = name
      self.dn = dn
      self.subtrees = {}
      self.leaves = []

   def __repr__(self, indent = 0):
      '''Returns an optionally indented, "friendly" representation of the
      (sub)tree on which this method is called.'''

      result = (' ' * indent) + self.dn + '\n'
      for leaf in self.leaves:
         result += leaf.__repr__(indent + 1)
      for subtree in self.subtrees:
         result += (' ' * (indent + 1)) + subtree + ' -> ' + \
            self.subtrees[subtree].__repr__(indent + 1).lstrip()
      return result

   def addSubtree(self, tree):
      '''Adds a subtree to the tree.'''
      self.subtrees[tree.name] = tree

   def addLeaf(self, leaf):
      '''Adds a leaf node (LDAPRecord) to the tree.'''
      self.leaves.append(leaf)

   def flatten(self, level = 0):
      '''Flattens the tree to a desired level (default 0). Flattening
      recursively moves all leaves below the flattening level to the
      flattening level. The flattening level is defined as the maximum
      subtree depth that will remain after the flattening operation. Thus,
      a flattening level of 0 will remove all subtrees, promoting all leaf
      nodes to be direct descendants of the root node. A flattening level
      of 1 allows 1 level of subtree, 2 permits 2 levels of subtree, and
      so forth. Flattening changes the tree in-place, and the changes
      cannot be undone (except by throwing away the tree and doing a new
      LDAP search).'''
      flatten_tree(self, level)

   def listGroups(self):
      '''Returns a list of all the "friendly" group names in the tree. Note
      that there can be duplicate group names, but each underlying tree node
      still has a unique DN corresponding to the ou represented by the
      group.'''
      return do_list_groups(self)

   def findGroup(self, group):
      '''Returns a list of tree nodes (LDAPTree type) that match the
      requested group name. Since group names are not unique, there can be
      multiple matches. If a specific group is desired, search by DN (which
      is unique) using the findDN method. Returns [] if the group is not
      found.'''
      return find_group(self, group)

   def findDN(self, dn):
      '''Returns the object (LDAPTree or LDAPRecord) with the requested DN.
      If no such object can be found, returns None. The DN may be an
      absolute DN (containing the base DN) or be relative to the base DN
      (ending with a comma).'''
      return opt_find_dn(self, full_dn(dn, self.dn))[0]

   def prune(self, dn):
      '''Removes an object (leaf or entire subtree) from the tree. Whether
      or not the garbage collector reclaims the removed object(s) depends
      upon the existence or non-existence of other references to the
      object(s). Does nothing if the requested DN is not in the tree. DN's
      may be absolute (containing the base DN) or relative to the base DN
      (ending with a comma).'''
      do_prune(self, full_dn(dn, self.dn))



class LDAPQuery(object):
   '''LDAPQuery objects represent a sequential search connection to an
   LDAP server. Multiple queries, with different base DN's, can be effected
   on a single connection. Each individual query will return its own
   LDAPTree object.'''

   def __init__(self, uri = 'ldap://localhost', starttls = False, \
                bind_dn = '', cred = ''):
      '''Construct the object and bind to an LDAP server at the specified
      uri. If starttls is requested, or an ldaps:// uri is used, the
      connection can be encrypted. If bind_dn and cred are not specified,
      an anonymous bind will be made to the server; otherwise, the provided
      DN and password will be used for the bind. To work with self-signed
      certificates, LDAPQuery objects will NOT perform certificate checking.'''

      self.uri = uri
      self.starttls = starttls
      self.bind_dn = bind_dn
      self.cred = cred
      self.server_handle = None

      # Do not check remote host certificate (to permit self-signed certs)
      ldap.set_option(ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_NEVER)
      self.server_handle = ldap.initialize(self.uri)
      if self.starttls:
         self.server_handle.start_tls_s()
      self.server_handle.simple_bind_s(self.bind_dn, self.cred)

   def __del__(self):
      '''The LDAPQuery object will remain bound to the server until the
      object is destroyed.'''

      self.server_handle.unbind_s()
      self.server_handle = None

   def query(self, base_dn = 'dc=localhost', limit = '(objectClass=*)'):
      '''Perform a sequential search on the LDAP server using the specified
      base_dn and return the results as an LDAPTree. The types of LDAP
      objects returned as LDAPRecords may be limited by using a limit
      (defaults to '(objectClass=*)'). See RFC 2254 for the search filter
      grammar (filters are referred to as limits, since filter is a reserved
      word in Python). Note that it is NOT necessary to include
      (objectClass=organizationalUnit) in the filter, since the ou hierarchy
      can be determined from the DN of each object.'''

      results = self.server_handle.search_s(base_dn, ldap.SCOPE_SUBTREE, limit)
      tree = results_to_tree(results, base_dn)
      return tree


##############################################################################

# Utility functions perform most of the actual work. These are implemented
# outside the actual class definitions for better clarity, especially since
# many of these functions are recursive.

# Strips the base DN off a provided DN (if the base DN is part of the
# input DN), so that subtrees of ou's are relative to the root (which
# corresponds to the base DN)
def strip_base(dn, base_dn):
   result = dn
   if dn.endswith(base_dn):
      result = dn[0:(-1 * len(base_dn))]
      if result.endswith(','):
         result = result[0:-1]
   return result.strip()

# Returns a tuple containing the outer-most group (ou segment without ou=)
# name in a string of ou's and the rest of the ou string.
#
# ASSUMPTION: ou's are only children of other ou's, or of the base DN. The
# tree functions will NOT work properly if ou's are children of other
# objects.
def get_group(oustr):
   parts = oustr.split(',')
   # magic here: the group name is actually the LAST ou, since the
   # hierarchy is reversed when presented as a string
   # (e.g. ou=flamejets,ou=DHCP,ou=Hosts)
   #
   # put_in_tree also needs to know the "rest" of the ou string -- and it
   # needs to be a string, so return a tuple
   return (parts[-1][3:], ','.join(parts[:-1]))

# Returns the list of ou's in a DN, starting at the first ou= segment in the
# DN. See the important ASSUMPTION about nested ou's, above.
def get_ou_list(dn):
   oulist = ''
   idx = dn.find('ou=')
   if idx >= 0:
      oulist = dn[idx:]
   return oulist

# Allow short DN's to be used with certain LDAPTree methods. DN's ending
# with a comma (,) are presumed to be relative to the base DN. Since most
# the tree operations expect a complete DN, this function permits the
# complete DN to be constructed from the short convenience "D"N.
def full_dn(in_dn, base_dn):
   out_dn = in_dn
   if in_dn.endswith(','):
      out_dn = in_dn + base_dn
   return out_dn

# Recursively find all group names defined in the tree (there may be
# duplicates)
def do_list_groups(tree):
   groups = tree.subtrees.keys()
   for subtree in tree.subtrees:
      groups.extend(do_list_groups(tree.subtrees[subtree]))
   return groups


###############################################################################
# Functions to convert raw LDAP search results into a tree.

# Recursive function to insert an object into the tree, creating subtrees
# as needed.
def put_in_tree(tree, oulist, dn, attrs, base_dn):
   if len(oulist) == 0:
      # Base case: insert the leaf node into the current tree level
      tree.addLeaf(LDAPRecord(dn, attrs))
   else:
      # Recursive case: find the subtree (group) if it exists, or recursively
      # create the subtree if it does not exist.
      subtree = None
      group, remainder = get_group(oulist)

      # To store the DN of the subtree (LDAP ou), re-create the RDN by
      # prepending 'ou=' to the group name. The DN can then be constructed
      # by appending the "base" DN -- which here is either the true search
      # base DN or the DN of the parent subtree.
      rdn = 'ou=' + group
      if group not in tree.subtrees:
         subtree = LDAPTree(group, rdn + ',' + base_dn)
         tree.addSubtree(subtree)
      else:
         subtree = tree.subtrees[group]
      put_in_tree(subtree, remainder, dn, attrs, rdn + ',' + base_dn)

# Iterative driver to run put_in_tree
def results_to_tree(results, base_dn):
   # The root node of the results tree uses the base DN for both its name
   # and DN.
   tree = LDAPTree(base_dn, base_dn)
   for result in results:
      dn, attrs = result
      # since the tree is organized by ou, it's a little silly to store
      # the ou's if the query didn't filter them out
      if not dn.startswith('ou='):
         ou_part = get_ou_list(strip_base(dn, base_dn))
         put_in_tree(tree, ou_part, dn, attrs, base_dn)
   return tree


##############################################################################

# Tree flattening routines. Trees are flattened in place to save memory.

def flatten_to_level(level, subtree):
   # Recursively flatten each subtree first
   for subsubtree in subtree.subtrees:
      flatten_to_level(level, subtree.subtrees[subsubtree])
   # Then "move" (copy) the leaves into the current level (which is a
   # subtree, or maybe the parent tree if flattening to zero). Don't
   # delete anything here to avoid iterator error. Also, we're merely
   # copying references, not a deep copy, so this is fairly quick.
   for leaf in subtree.leaves:
      level.addLeaf(leaf)

def flatten_tree(tree, flatten):
   # MUST iterate on the .keys() (a list containing a COPY of the dict keys)
   # here, or an exception will occur because the dict would change during
   # iteration
   for subtree in tree.subtrees.keys():
      # Recursively do nothing until desired level is reached
      if flatten > 0:
         flatten_tree(tree.subtrees[subtree], flatten - 1)
      else:
         flatten_to_level(tree, tree.subtrees[subtree])
         del tree.subtrees[subtree]


###############################################################################

# Search routines

# Returns a LIST of trees, empty list if no matching group found
def find_group(tree, group):
   retlist = []
   for subtree in tree.subtrees:
      retlist.extend(find_group(tree.subtrees[subtree], group))
      if subtree == group:
         retlist.append(tree.subtrees[subtree])
   return retlist


# Returns a tuple (leaf_or_subtree, parent) if the requested DN can be
# found in the tree. Returns None if the DN is not found.
def opt_find_dn_recurse(tree, ou_str, dn):
   if ou_str == '':
      for leaf in tree.leaves:
         if leaf.dn == dn:
            return (leaf, tree)
   else:
      group, remainder = get_group(ou_str)
      subtree = tree.subtrees[group]
      # What if user requested an ou?
      if subtree.dn == dn:
         return (subtree, tree)
      return opt_find_dn_recurse(subtree, remainder, dn)
   return None

# Driver for the recursive DN search algorithm. See important ASSUMPTION
# about ou nesting near get_group function above.
def opt_find_dn(tree, dn):
   base_dn = tree.dn
   if not dn.endswith(base_dn):     # requested DN not possibly in tree
      return None
   if dn == base_dn:
      return (tree, None)
   rdn = strip_base(dn, base_dn)
   ou_list = get_ou_list(rdn)
   return opt_find_dn_recurse(tree, ou_list, dn)


###############################################################################

# Tree pruning: remove a leaf or entire branch (subtree)
def do_prune(tree, dn):
   location, parent = opt_find_dn(tree, dn)
   if type(location) is LDAPTree:
      if parent:
         del parent.subtrees[location.name]
      # if parent is None then user is trying to prune the whole tree
      # fail silently: the user should just throw the tree away
   elif type(location) is LDAPRecord:
      parent.leaves.remove(location)
   # else do nothing: dn not in tree


###############################################################################

# Test code: specific to our test system
if __name__ == '__main__':
   query = LDAPQuery('ldap://firebox.cs.clemson.edu', True)
   tree = query.query('ou=Hosts,dc=firebox,dc=cs,dc=clemson,dc=edu')
   #print tree
   #tree.flatten()
   #print tree
   #flamejets = tree.findGroup('flamejets')
   #print flamejets
   #nogroup = tree.findGroup('silly')
   #print nogroup
   #oiltanks = tree.findDN('ou=oiltanks,ou=DHCP,')
   #print oiltanks
   #nodn = tree.findDN('cn=silly,ou=nothing,dc=firebox,dc=cs,dc=clemson,dc=edu')
   #print nodn
   #tree.prune('cn=oiltank1,ou=oiltanks,ou=DHCP,')
   #oiltanks = tree.findDN('ou=oiltanks,ou=DHCP,')
   #print oiltanks
   #tree.prune('ou=oiltanks,ou=DHCP,')
   #print tree
   tree.flatten(1)
   print tree.listGroups()
